package com.wjf.wjfaiagent.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.stereotype.Component;

@Component
public class QueryRewriter {
    private final RewriteQueryTransformer queryTransformer;
    public QueryRewriter(ChatModel dashscopeChatModel){
        RewriteQueryTransformer queryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(ChatClient.builder(dashscopeChatModel))
                .build();
        this.queryTransformer = queryTransformer;
    }
    public String doQueryRewrite(String prompt){
        Query query = queryTransformer.transform(new Query(prompt));
        return query.text();
    }
}
